package com.baidu.paddle.lite.demo.segmentation

import android.content.Context
import android.graphics.Bitmap
import android.util.Log
import com.baidu.paddle.lite.demo.segmentation.config.Config
import com.baidu.paddle.lite.demo.segmentation.preprocess.Preprocess
import com.baidu.paddle.lite.demo.segmentation.visual.Visualize
import com.zhaiker.paddlesegmentation.R
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext

object PaddleSegHelper {

    // model config
    private var config: Config = Config()

    private var predictor: Predictor = Predictor()

    private var preprocess: Preprocess = Preprocess()

    private var visualize: Visualize = Visualize()

    private lateinit var mContext : Context

    private fun initSync(context: Context){
        mContext = context
        val model_path: String = mContext.getString(R.string.MODEL_PATH_DEFAULT)
        val label_path : String = mContext.getString(R.string.LABEL_PATH_DEFAULT)
        val image_path : String = ""

        val cpu_thread_num: Int = mContext.getString(R.string.CPU_THREAD_NUM_DEFAULT).toInt()
        val cpu_power_mode: String = mContext.getString(R.string.CPU_POWER_MODE_DEFAULT)

        val input_color_format: String = mContext.getString(R.string.INPUT_COLOR_FORMAT_DEFAULT)

        val input_shape =
            Utils.parseLongsFromString(
                mContext.getString(R.string.INPUT_SHAPE_DEFAULT), ","
            )

        config.init(
            model_path, label_path, image_path, cpu_thread_num, cpu_power_mode,
            input_color_format, input_shape
        )
        preprocess.init(config)

        predictor.init(mContext, config)
    }

    suspend fun init(context: Context) = withContext(Dispatchers.IO){
        initSync(context)
    }

    private fun segImageSync(input : Bitmap?) : Bitmap?{
        if (input != null && predictor.isLoaded()) {
            predictor.setInputImage(input)

            if(predictor.isLoaded() && predictor.runModel(preprocess, visualize)){
                Log.d("PaddleSegmentation","Inference time: " + predictor.inferenceTime() + " ms")
                val outputImage = predictor.outputImage()
                if (outputImage != null) {
                    return outputImage
                }
                Log.d("PaddleSegmentation",predictor.outputResult())
            }else{
                Log.d("PaddleSegmentation","分割失败")
            }
        }
        return null
    }

    suspend fun segImage(input : Bitmap?) : Bitmap? = withContext(Dispatchers.Default){
        return@withContext segImageSync(input)
    }

}